[
    "gather",
    "reset_trace_state",
    "ihfft",
    "lower_parallel_callable",
    "DimensionHandlerTracer",
    "register_keypaths",
    "associative_scan",
    "after_all",
    "isf",
    "recipe_to_eqn",
    "xavier_uniform",
    "unregister_vmappable",
    "donation_vector",
    "identity",
    "global_aval_to_result_handler",
    "AbstractValue",
    "annotate_function",
    "tree_all",
    "arg_info_all",
    "shape_tensor",
    "sum_shapes",
    "StoreException",
    "object_",
    "trace_to_subjaxpr_dynamic2",
    "register_lowering",
    "blackman",
    "DynamicJaxprTracer",
    "pdot",
    "dedup_referents",
    "cdouble",
    "DynamicJaxprTrace",
    "fftn",
    "Tracer",
    "axis_index",
    "naryop_dtype_rule",
    "igamma",
    "mul",
    "PRNGKey",
    "is_constant_dim",
    "sign",
    "where",
    "zeros_like_jaxval",
    "is_op_sharding_replicated",
    "ThreadLocalState",
    "_deprecated_DeviceArray",
    "lower_jaxpr_to_fun",
    "dynamic_update_slice",
    "dynamic_slice",
    "_get_sharding_specs",
    "tree_reduce",
    "lpmn_values",
    "multivariate_normal",
    "custom_vjp",
    "broadcast_batcher",
    "Var",
    "eval_jaxpr",
    "MapTrace",
    "normal",
    "t",
    "_deprecated_OpShardingSharding",
    "MapTracer",
    "__getattr__",
    "ShardedAxis",
    "register_constant_handler",
    "select_n",
    "log1p",
    "outfeed",
    "put",
    "token_type",
    "call_transpose",
    "segment_sum",
    "new_jaxpr_eqn",
    "get_printoptions",
    "partitioned_sharding_spec",
    "vsplit",
    "canonicalize_shape",
    "constant",
    "diag_indices",
    "fori_loop",
    "can_cast",
    "axis_frame",
    "floating",
    "_deprecated_make_sharded_device_array",
    "to_elt",
    "det",
    "weibull_min",
    "log",
    "Replicated",
    "vstack",
    "roll",
    "tril_indices",
    "zeros_like_array",
    "new_sublevel",
    "xla_pmap_impl_lazy",
    "PmapComputation",
    "merge_linear_aux",
    "UnloadedPmapExecutable",
    "_deprecated_ShapedArray",
    "DebugEffect",
    "PileAxis",
    "unop",
    "Backend",
    "UnloadedMeshExecutable",
    "subst_axis_names",
    "scale_and_translate",
    "abstractify",
    "conv_transpose",
    "inv",
    "_PRNGKeyArray",
    "lstsq",
    "treedef_tuple",
    "tridiagonal_solve",
    "LambdaBinding",
    "solve_triangular",
    "print_saved_residuals",
    "RegularGridInterpolator",
    "sharding_spec_sharding_proto",
    "cho_factor",
    "UnexpectedTracerError",
    "ones",
    "ModuleContext",
    "xla_destructure",
    "std",
    "ConcretizationTypeError",
    "defjvp2",
    "trace",
    "schur",
    "xeinsum",
    "_deprecated_curry",
    "indices",
    "inexact",
    "clamp",
    "dirichlet",
    "kaiming_uniform",
    "infer_lambda_input_type",
    "int_",
    "process_index",
    "conv_general_dilated",
    "instantiate_const_at",
    "pareto",
    "UndefinedPrimal",
    "diff_dim",
    "dot",
    "segment_max",
    "Trace",
    "argnums_partial",
    "convert_element_type",
    "in1d",
    "_get_and_check_device_assignment",
    "block_until_ready",
    "get_compile_options",
    "Literal",
    "total_ordering",
    "eigh",
    "batch_jaxpr",
    "bernoulli",
    "ShardingSpec",
    "defreducer",
    "JaxprTrace",
    "device_put_replicated",
    "erfinv",
    "conj",
    "lpmn",
    "local_device_count",
    "namedtuple",
    "dense_int_elements",
    "xlog1py",
    "ComplexWarning",
    "RoundingMethod",
    "broadcast_to_rank",
    "abstract_eval_fun",
    "betainc",
    "flatnonzero",
    "move_binders_to_front",
    "BoundedAxisSize",
    "TraceState",
    "get_primitive_transpose",
    "pswapaxes",
    "ExecuteReplicated",
    "tuple_insert",
    "igammac",
    "array_equiv",
    "defbilinear",
    "call_padding_rule",
    "DropVar",
    "gaussian_kde",
    "parallel_callable",
    "hessian",
    "xla_pmap_impl",
    "MeshComputation",
    "piecewise",
    "used_axis_names",
    "host_id",
    "lt",
    "callback",
    "bincount",
    "traverse_jaxpr_params",
    "ndim",
    "take",
    "replicate",
    "float16",
    "extend_axis_env_nd",
    "conv_general_dilated_patches",
    "ParallelCallableInfo",
    "complex64",
    "dynamic_slice_in_dim",
    "ge",
    "tree_unflatten",
    "valid_jaxtype",
    "Lowered",
    "eye",
    "xlogy",
    "scatter_apply",
    "xla_computation_to_mlir_module",
    "iterable",
    "ppermute",
    "xavier_normal",
    "device_get",
    "op_sharding_to_indices",
    "pp_kv_pairs",
    "cg",
    "set_printoptions",
    "sometrue",
    "defvectorized",
    "TileVectorize",
    "float64",
    "gumbel",
    "householder_product",
    "Jaxpr",
    "dynamic_index_in_dim",
    "iscomplexobj",
    "hsplit",
    "call_impl",
    "start_server",
    "jvp_jaxpr",
    "polydiv",
    "InputsHandler",
    "get_local_aval",
    "resize",
    "all_to_all",
    "reassemble_concat_axes",
    "dtype",
    "lecun_normal",
    "DictKey",
    "ClosedJaxpr",
    "expand_dims",
    "i0",
    "partial_eval_jaxpr_custom_rule_not_implemented",
    "laplace",
    "histogram",
    "ConcreteArray",
    "keystr",
    "uniform",
    "make_jaxpr_effects",
    "histogramdd",
    "reduce_window_shape_tuple",
    "vectorize",
    "load",
    "convert_envvars_to_constvars",
    "closure_convert",
    "record_event_duration_secs",
    "cholesky",
    "segment_min",
    "partial_eval_jaxpr_nounits",
    "register_translation",
    "amin",
    "split_dict",
    "eval_context",
    "DimExprEvaluator",
    "greater_equal_shape",
    "cumprod",
    "wraps",
    "deflinear2",
    "index_in_dim",
    "is_finite",
    "padtype_to_pads",
    "_version_as_tuple",
    "float_",
    "FlattenedIndexKey",
    "map",
    "_initial_style_jaxpr",
    "acosh",
    "CustomVJPException",
    "join_named_shapes",
    "sub",
    "rfftn",
    "acos",
    "PartitionSpec",
    "all",
    "_deprecated_Device",
    "JaxprEqnRecipe",
    "register_pytree_with_keys",
    "resource_typecheck",
    "closed_backward_pass",
    "trace_to_jaxpr",
    "empty_like",
    "symbolic_equal_one_of_dim",
    "union1d",
    "tril",
    "check_backend_matches",
    "grad",
    "unpair_pval",
    "lu_pivots_to_permutation",
    "find_top_trace",
    "broadcasted_iota",
    "register_pytree_node_class",
    "map_transpose",
    "TranslationRule",
    "array_equal",
    "value_and_grad",
    "check_type",
    "key_data",
    "make_device_array",
    "bartlett",
    "delta_orthogonal",
    "sf",
    "float8_e4m3fn",
    "truncated_normal",
    "jaxprs_in_params",
    "psum_scatter",
    "pmin",
    "KeyArray",
    "stop_gradient",
    "TraceAnnotation",
    "wrap_init",
    "initialize",
    "igamma_grad_a",
    "mean",
    "clz",
    "full_lower",
    "vtile",
    "add",
    "untile_aval_nd",
    "pow",
    "bessel_i1e",
    "conv_general_dilated_local",
    "host_ids",
    "permutation",
    "visualize_array_sharding",
    "betaln",
    "new_eqn_recipe",
    "dsplit",
    "dstack",
    "is_undefined_primal",
    "fold_in",
    "unsafe_rbg_key",
    "full_like",
    "unzip2",
    "dce_jaxpr",
    "rng_uniform",
    "trace_to_jaxpr_dynamic2",
    "stop_server",
    "one_hot",
    "typematch",
    "rfft",
    "transpose",
    "get_global_aval",
    "wrap_name",
    "rev",
    "empty",
    "vmap",
    "zeros_like",
    "size",
    "map_coordinates",
    "exp1",
    "broadcast_to",
    "ndtri",
    "ball",
    "fromfile",
    "dct",
    "logsf",
    "TileManual",
    "reduce",
    "save",
    "copy",
    "is_constant_shape",
    "do_subst_axis_names_jaxpr",
    "concrete_aval",
    "ref",
    "dense_bool_elements",
    "dynamic_update_slice_in_dim",
    "ResultsHandler",
    "gammaln",
    "shaped_abstractify",
    "Sublevel",
    "EvalTrace",
    "logpdf",
    "cache",
    "extend_jaxpr_stack",
    "local_aval_to_result_handler",
    "greater_equal_dim",
    "conv_transpose_shape_tuple",
    "number",
    "stride_shape",
    "replace_float0s",
    "_call_lowering",
    "fft2",
    "save_device_memory_profile",
    "jaxpr_subcomp",
    "Zero",
    "map_bind_with_continuation",
    "scan",
    "rbg_key",
    "tree_leaves",
    "pp_jaxpr",
    "concatenate",
    "defbroadcasting",
    "AttributeKeyPathEntry",
    "dimension_as_value",
    "str_eqn_compact",
    "meshgrid",
    "float32",
    "subst_axis_names_eqn",
    "standard_jvp",
    "broadcast_in_dim",
    "apply_todos",
    "detrend",
    "JaxprStackFrame",
    "spec_to_indices",
    "tree_flatten_with_path",
    "is_opaque_dtype",
    "standard_primitive",
    "cur_sublevel",
    "rsqrt",
    "cdf",
    "erf",
    "tracers_to_jaxpr",
    "generic",
    "dynamic_update_index_in_dim",
    "ir_constants",
    "Mesh",
    "tridiagonal",
    "ifftn",
    "linear_call",
    "dce_jaxpr_consts",
    "signedinteger",
    "treedef_children",
    "i1e",
    "cummin",
    "NonConcreteBooleanIndexError",
    "ptp",
    "SequenceKey",
    "GSPMDSharding",
    "dce_jaxpr_call_rule",
    "SymbolicZero",
    "make_array_from_callback",
    "TokenSet",
    "bitwise_and",
    "softmax",
    "convolve2d",
    "int16",
    "tree_transpose",
    "RandomAlgorithm",
    "aval_to_ir_type",
    "pshuffle",
    "issubdtype",
    "pp_jaxpr_skeleton",
    "issubsctype",
    "gt",
    "ResizeMethod",
    "tile",
    "ceil",
    "new_base_main",
    "ConcatAxis",
    "int8",
    "map_bind",
    "iota",
    "tan",
    "fromfunction",
    "bitwise_xor",
    "mask_indices",
    "pad",
    "triu",
    "i0e",
    "Pile",
    "multigammaln",
    "custom_jvp",
    "host_count",
    "switch",
    "all_gather",
    "lower_jaxpr_to_module",
    "delete",
    "triu_indices_from",
    "tree_structure",
    "sum_dim",
    "trace_to_jaxpr_nounits",
    "logcdf",
    "setxor1d",
    "kaiser",
    "_sum_tangents",
    "square",
    "dilate_shape",
    "pad_jaxpr",
    "custom_vmap",
    "character",
    "mesh_sharding_specs",
    "Compiled",
    "hanning",
    "round",
    "InDBIdx",
    "hstack",
    "jacfwd",
    "subst_axis_names_var",
    "tri",
    "live_arrays",
    "aval_to_ir_types",
    "Shard",
    "batch_jaxpr_axes",
    "check_eqn",
    "PileTy",
    "tree_map",
    "sqrt",
    "SingleDeviceSharding",
    "fftshift",
    "new_mesh_sharding_specs",
    "logspace",
    "xla_computation",
    "printoptions",
    "Partial",
    "fftfreq",
    "_deprecated_device_put",
    "tril_indices_from",
    "get_referent",
    "nanargmin",
    "PyTreeDef",
    "typecompat",
    "slice",
    "diagflat",
    "sig_info",
    "new_main",
    "cummax",
    "shift_right_arithmetic",
    "recast_to_float0",
    "shard_arg",
    "CustomJVPException",
    "_deprecated_Mesh",
    "ifft2",
    "wald",
    "trim_zeros",
    "AxisEnv",
    "_deprecated_PartitionSpec",
    "ne",
    "parameter",
    "select",
    "cumlogsumexp",
    "pure_callback",
    "find_replicas",
    "approx_min_k",
    "named_scope",
    "jacobian",
    "pp_eqns",
    "PartialVal",
    "top_k",
    "lecun_uniform",
    "jit",
    "vtile_by_mesh",
    "custom_transpose",
    "are_op_shardings_equal",
    "generalized_normal",
    "logpmf",
    "pmap",
    "sort_key_val",
    "closed_call_partial_eval_custom_rule",
    "default_backend",
    "lower_mesh_computation",
    "process_count",
    "trace_to_subjaxpr_dynamic",
    "disable_jit",
    "array_repr",
    "segment_prod",
    "fliplr",
    "population_count",
    "_create_pmap_sharding_spec",
    "EmapInfo",
    "ones_like",
    "stage_parallel_callable",
    "mapped_aval",
    "scatter",
    "atan",
    "polygamma",
    "while_loop",
    "dot_general",
    "eq",
    "NotMapped",
    "register_pytree_node",
    "Token",
    "infeed",
    "GetAttrKey",
    "float8_e5m2",
    "finfo",
    "erfc",
    "partition_pvals",
    "symbolic_equal_shape",
    "Sharding",
    "ArgInfo",
    "Effect",
    "Buffer",
    "Array",
    "argmin",
    "apply_over_axes",
    "substitute_vars_in_output_ty",
    "register_vmappable",
    "abs",
    "squeeze",
    "real",
    "cos",
    "device_memory_profile",
    "rayleigh",
    "nonzero",
    "glorot_uniform",
    "Device",
    "JaxprPpContext",
    "remat",
    "ndarray",
    "print",
    "uint32",
    "named_call",
    "le",
    "linear_transpose",
    "same_referent",
    "istft",
    "fromiter",
    "apply_along_axis",
    "ix_",
    "create_token",
    "savez",
    "WeakRefList",
    "pp_kv_pair",
    "vjp",
    "backward_pass",
    "bicgstab",
    "gensym",
    "leaked_tracer_error",
    "geomspace",
    "moveaxis",
    "he_uniform",
    "integer",
    "neg",
    "cho_solve",
    "JAXTypeError",
    "defjvp_zero",
    "psum",
    "canonicalize_dtype",
    "partial_eval_jaxpr_custom",
    "escaped_tracer_error",
    "make_ir_context",
    "rfftfreq",
    "tensordot",
    "breakpoint",
    "nancumprod",
    "_CppDeviceArray",
    "sph_harm",
    "JaxprPpSettings",
    "BatchTrace",
    "cauchy",
    "isrealobj",
    "_DeviceArray",
    "rng_bit_generator",
    "typecheck",
    "shift_right_logical",
    "convolve",
    "is_empty_shape",
    "tuple_delete",
    "OutDBIdx",
    "double",
    "ConstVar",
    "DynamicAxisEnv",
    "correlate",
    "array_str",
    "tensorsolve",
    "multi_dot",
    "uint16",
    "devices",
    "jaxpr_uses_outfeed",
    "trace_to_jaxpr_final",
    "LoweringRuleContext",
    "bilinear_transpose",
    "fromstring",
    "from_elt",
    "ppf",
    "cbrt",
    "sum",
    "flip",
    "promote_types",
    "orthogonal",
    "PmapExecutable",
    "register_pytree_with_keys_class",
    "nextafter",
    "correlate2d",
    "eig",
    "expm1",
    "from_dlpack",
    "symbolic_equal_dim",
    "HashableFunction",
    "ZeroIfMapped",
    "stash_axis_env",
    "print_environment_info",
    "unpack_concat_axes",
    "make_jaxpr",
    "product",
    "with_sharding_constraint",
    "uint8",
    "complex_",
    "JaxprTypeError",
    "TracerArrayConversionError",
    "checkpoint",
    "div",
    "extract",
    "unsignedinteger",
    "threefry2x32_key",
    "LoweringResult",
    "shuffle",
    "bitwise_or",
    "GatherDimensionNumbers",
    "is_special_dim_size",
    "JAXIndexError",
    "conv_general_permutations",
    "lower_fun",
    "as_named_shape",
    "choice",
    "partial",
    "ArrayMapping",
    "StepTraceAnnotation",
    "Wrapped",
    "int32",
    "rademacher",
    "register_event_duration_secs_listener",
    "polymul",
    "pmean",
    "register_initial_style_primitive",
    "approx_max_k",
    "as_hashable_function",
    "entr",
    "matchaxis",
    "unmapped_aval",
    "SPMDBatchTrace",
    "hessenberg",
    "svd",
    "call_bind_with_continuation",
    "NoSharding",
    "Chunked",
    "emit_python_callback",
    "var",
    "bdim_at_front",
    "XLACompatibleSharding",
    "GatherScatterMode",
    "build_tree",
    "broadcast",
    "array",
    "GetitemKeyPathEntry",
    "result_type",
    "add_batched",
    "diff_shape",
    "is_dim",
    "compress",
    "NamedSharding",
    "JaxprEqn",
    "dce_jaxpr_closed_call_rule",
    "DebugInfo",
    "concretization_function_error",
    "result_info",
    "sinh",
    "linear_jvp",
    "check_jaxpr",
    "flatten_lowering_ir_args",
    "conv",
    "row_stack",
    "make_iota",
    "csingle",
    "irfft",
    "effects_barrier",
    "conv_with_general_padding",
    "SPMDAxisContext",
    "reduce_window",
    "i64_attr",
    "ifftshift",
    "Initializer",
    "OrderedDictType",
    "safe_zip",
    "alltrue",
    "logsumexp",
    "pdf",
    "welch",
    "argmax",
    "pp_vars",
    "chisquare",
    "scatter_add",
    "batch_matmul",
    "shard_aval",
    "core_call_lowering",
    "threefry_2x32",
    "exp",
    "zeros",
    "pp_eqn",
    "array_split",
    "subst_axis_names_jaxpr",
    "naryop",
    "scatter_mul",
    "transfer_guard",
    "imag",
    "register_event_listener",
    "merge_mlir_modules",
    "sort",
    "gamma",
    "linearize",
    "spence",
    "TracerId",
    "conv_shape_tuple",
    "unop_dtype_rule",
    "expit",
    "debug_info_final",
    "batch_custom_vjp_bwd",
    "zero_jvp",
    "column_stack",
    "cumproduct",
    "complex",
    "kaiming_normal",
    "exponential",
    "jvp",
    "csd",
    "raise_as_much_as_possible",
    "uint64",
    "MapPrimitive",
    "linspace",
    "IndexedAxisSize",
    "rfft2",
    "intersect1d",
    "is_vmappable",
    "bessel_i0e",
    "scatter_max",
    "isin",
    "logistic",
    "ShardInfo",
    "_zeros_like_pytree",
    "tanh",
    "array_mapping_to_axis_resources",
    "bfloat16",
    "diag_indices_from",
    "TracerAsName",
    "min",
    "register_collective_primitive",
    "toeplitz",
    "trace_state_clean",
    "TraceStack",
    "msort",
    "backend_compile",
    "make_array_from_single_device_arrays",
    "split_list",
    "nancumsum",
    "start_trace",
    "irfftn",
    "asarray",
    "_pmap_sharding_spec",
    "int64",
    "tile_aval_nd",
    "broadcast_arrays",
    "nanargmax",
    "call_partial_eval_custom_rule",
    "reducer_batcher",
    "shift_left",
    "he_normal",
    "debug_info",
    "zeros_like_batched",
    "average",
    "CallPrimitive",
    "default_prng_impl",
    "NameGatheringSubst",
    "module_to_bytecode",
    "scalar_type_of",
    "gammaincc",
    "local_devices",
    "argwhere",
    "setdiff1d",
    "sequential_vmap",
    "find_partitions",
    "cosh",
    "sin",
    "lu",
    "solve",
    "i1",
    "power",
    "PartitionInfo",
    "isscalar",
    "device_put",
    "ndtr",
    "conv_general_shape_tuple",
    "def_trivial_padding",
    "defjvp",
    "maybe_find_leaked_tracers",
    "ScatterDimensionNumbers",
    "vq",
    "all_leaves",
    "funm",
    "JVPTrace",
    "module_to_string",
    "ensure_compile_time_eval",
    "fft",
    "_one_replica_buffer_indices",
    "eval_shape",
    "instantiate_zeros_aval",
    "_is_unspecified",
    "pp_jaxprs",
    "maybe_extend_axis_env",
    "stop_trace",
    "move_binders_to_back",
    "global_avals_to_results_handler",
    "inspect_array_sharding",
    "integer_pow",
    "custom_gradient",
    "pmf",
    "categorical",
    "f",
    "vectorized_batcher",
    "tensorinv",
    "complex128",
    "randint",
    "subvals",
    "sharding_to_proto",
    "pp_aval",
    "scatter_min",
    "einsum",
    "conv_dimension_numbers",
    "cond",
    "triu_indices",
    "record_event",
    "stride_dim",
    "insert",
    "visualize_sharding",
    "slice_in_dim",
    "join_effects",
    "custom_root",
    "WrappedFun",
    "scan_bind",
    "repeat",
    "index_take",
    "custom_linear_solve",
    "zeta",
    "dilate_dim",
    "get_aval",
    "MainTrace",
    "JaxprTracer",
    "ShapedArray",
    "isposinf",
    "gelu",
    "rearrange_binders",
    "lattice_join",
    "DShapedArray",
    "complexfloating",
    "check_valid_jaxtype",
    "bitwise_not",
    "AbstractToken",
    "standard_jvp2",
    "split_merge",
    "DynamicAxisEnvFrame",
    "atan2",
    "bool_",
    "ReplicaInfo",
    "has_opaque_dtype",
    "i32_attr",
    "instantiate_zeros",
    "unique",
    "ifft",
    "reciprocal",
    "normalize",
    "FreeVar",
    "batch",
    "concrete_or_error",
    "same_shape_sizes",
    "jacrev",
    "shard_args",
    "ravel_multi_index",
    "maxwell",
    "MeshExecutable",
    "NamedShape",
    "InconclusiveDimensionOperation",
    "trivial_ctx",
    "loggamma",
    "histogram2d",
    "max",
    "beta",
    "frombuffer",
    "any",
    "hamming",
    "toposort",
    "uint",
    "get_num_partitions",
    "device_put_sharded",
    "glorot_normal",
    "clear_backends",
    "add_jaxvals",
    "tree_flatten",
    "BatchTracer",
    "LoweringRule",
    "JVPTracer",
    "ConvDimensionNumbers",
    "stack",
    "safe_map",
    "tie_in",
    "cumsum",
    "device_count",
    "flipud",
    "Unstacked",
    "deflinear",
    "reshape",
    "tree_map_with_path",
    "rebase_donate_argnums",
    "asin",
    "ShardingContext",
    "isneginf",
    "extend_axis_env",
    "unravel_index",
    "batch_jaxpr2",
    "TranslationContext",
    "histogram_bin_edges",
    "linear_transpose2",
    "lower_sharding_computation",
    "axis_groups",
    "primitive_uses_outfeed",
    "hfft",
    "irfft2",
    "choose",
    "double_sided_maxwell",
    "UnshapedArray",
    "variance_scaling",
    "trace_to_jaxpr_dynamic",
    "rem",
    "Precision",
    "PmapSharding",
    "qr",
    "Primitive",
    "_deprecated_ShardedDeviceArray",
    "subjaxprs",
    "atanh",
    "treedef_is_leaf",
    "stft",
    "sqrtm",
    "reconcile_num_partitions",
    "raise_to_shaped",
    "ReplicaAxisContext",
    "checkpoint_name",
    "split",
    "broadcast_shapes",
    "standard_abstract_eval",
    "shutdown",
    "zeros_like_aval",
    "bitcast_convert_type",
    "flatten_axes",
    "primal_dtype_to_tangent_dtype",
    "gammainc",
    "ShapeDtypeStruct",
    "place",
    "single",
    "pmax",
    "apply_primitive",
    "digamma",
    "amax",
    "TracerIntegerConversionError",
    "triangular_solve",
    "erf_inv",
    "pp_var",
    "canonicalize_dim",
    "diag",
    "reduce_precision",
    "zip",
    "floor",
    "gmres",
    "poisson",
    "iinfo",
    "XlaOp",
    "PositionalSharding",
    "add_tangents",
    "arange",
    "dctn",
    "trace_to_jaxpr_final2",
    "local_avals_to_results_handler",
    "divide_shape_sizes",
    "asinh",
    "ConstantHandler",
    "dtype_to_ir_type",
    "full",
    "lgamma",
    "einsum_path",
    "collapse",
    "pp_jaxpr_eqn_range",
    "prod",
    "roots",
    "flexible",
    "random_gamma_grad",
    "ir_constant",
    "shape",
    "PRNGKeyArray"
]